import  dotenv
import weaviate
from langchain.retrievers import ContextualCompressionRetriever
from langchain.retrievers.document_compressors import CohereRerank
from langchain_openai import OpenAIEmbeddings
from langchain_weaviate import WeaviateVectorStore
from weaviate.auth import AuthApiKey

import os

dotenv.load_dotenv()

# 创建数据库
embedding = OpenAIEmbeddings(model='text-embedding-3-small')

client = weaviate.connect_to_weaviate_cloud(
    skip_init_checks=True,
    cluster_url=os.getenv("WAEVIATE_URL"),
    auth_credentials=AuthApiKey(os.getenv("WEAVIATE_KEY"))
)

db = WeaviateVectorStore(
    client=client,
    index_name="TestParent",
    text_key="text",
    embedding = OpenAIEmbeddings(model="text-embedding-3-small")
)

rerank = CohereRerank(model="rerank-multilingual-v3.0")
# 构建检索器
retriever = ContextualCompressionRetriever(
    base_retriever = db.as_retriever(),
    base_compressor = rerank

)

#检索并返回
search_docs =retriever.invoke("分享关于LLMops 的一些应用配置")
print(search_docs)
for i in search_docs:
    print(i)
print(len(search_docs))